#!/usr/local/bin/python

"""
 Split input fastq files in files of n-million reads each and execute bwa on each of them.
"""
import os
import subprocess
import sys
import re

## bsub -R "rusage[mem=4096]" -n 1 -o log.txt < bwa_align.py


# ----------------------------[ Settings ]--------------------------------------

GENOME= '/lustre/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa'
FQ_DIR= 'fastq_test'      ## All the files in this dir will be processed.
FASTQ_EXT= '\.fq\.gz$' ## Regex to get fastq files. This match will be stripped off the name and the rest used for the output files. MEMO: regex '\.s_[1-8]_sequence\.txt\.gz$' to strip full Solexa name from suffix.
BWA_ALN_OPT= ''
PREFIXTAG= 'splitfile_' ## This tag is to recognize files that come from splitting 

genome_dict= { '/lustre/reference_data/genomes/Mus_musculus_NCBI_v37/mmu.fa': 'mm9',
               '/lustre/reference_data/genomes/Homo_sapiens_NCBI_v36.3/hsa.fa': 'hg18'} ## Tag for output files to look like '/lustre/.../fastq/bham-470.bwa.mm9.bam'

# ------------------------------------------------------------------------------

if not os.path.exists('splitfiles'):
    os.makedirs('splitfiles')

fastq_files= os.listdir(FQ_DIR)
genome_tag= genome_dict[GENOME]

n= 1
for fq in fastq_files:
    basename= re.sub(FASTQ_EXT, '', fq)
    fq_path= os.path.join(FQ_DIR, fq)
    fq_split= os.path.join('splitfiles', PREFIXTAG + fq) 
#    sai= os.path.join(OUTPUT_DIR, fq + '.bwa.' + genome_tag + '.sai')
##    sam= os.path.join(OUTPUT_DIR, re.sub(FASTQ_EXT, '.bwa.%s.sam' %genome_tag, fq) )
#    bam= os.path.join(OUTPUT_DIR, re.sub(FASTQ_EXT, '.bwa.%s.unsorted.bam' %genome_tag, fq) )
#    bam_sorted= os.path.join(OUTPUT_DIR, re.sub(FASTQ_EXT, '.bwa.%s' %genome_tag, fq) )
    cmd_file= os.path.join('splitfiles', fq + '.sh')
    shfile= open(cmd_file, 'w')
    cmd= """
#!/bin/bash

set -e

## rm %(fq_split)s*    
zcat %(fq_path)s | split -d -a 3 -l 4000 - %(fq_split)s
for fq in `ls %(fq_split)s*`
    do
    echo '#!/bin/bash' > $fq.sh
    echo "set -e" >> $fq.sh
    echo "/home/brown22/local/bin/bwa aln %(bwa_aln_opt)s %(genome)s $fq > $fq.sai" >> $fq.sh
    echo "/home/brown22/local/bin/bwa samse %(genome)s $fq.sai $fq > $fq.sam" >> $fq.sh
    echo "/home/berald01/applications/samtools/samtools view -bS $fq.sam > $fq.bam" >> $fq.sh
    echo "rm $fq.sai; rm $fq.sam; rm $fq" >> $fq.sh
    bsub -q test -R "rusage[mem=4096]" < $fq.sh
    done    
         """ %{'fq_path':fq_path, 'fq_split':fq_split , 'bwa_aln_opt':BWA_ALN_OPT, 'genome':GENOME}
    shfile.write(cmd)
    shfile.close()
    os.chmod(cmd_file, 0744)
    p= subprocess.Popen('bsub -J %s -q test -R "rusage[mem=1024]" < %s' %('j-' + str(n).zfill(3), cmd_file), shell= True)
    ## p= subprocess.Popen('bsub -w done(fq*) -q test -R "rusage[mem=1024]" < %s' %(fq, cmd_file), shell= True)
    n += 1
sys.exit()
